{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "# os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "import tensorflow as tf\n",
    "tf.compat.v1.enable_eager_execution()\n",
    "\n",
    "import numpy as np\n",
    "import imageio\n",
    "import json\n",
    "import random\n",
    "import time\n",
    "import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import run_nerf\n",
    "\n",
    "from load_llff import load_llff_data\n",
    "from load_deepvoxels import load_dv_data\n",
    "from load_blender import load_blender_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "basedir = './logs'\n",
    "expname = 'fern_example'\n",
    "\n",
    "config = os.path.join(basedir, expname, 'config.txt')\n",
    "print('Args:')\n",
    "print(open(config, 'r').read())\n",
    "parser = run_nerf.config_parser()\n",
    "\n",
    "args = parser.parse_args('--config {} --ft_path {}'.format(config, os.path.join(basedir, expname, 'model_200000.npy')))\n",
    "print('loaded args')\n",
    "\n",
    "images, poses, bds, render_poses, i_test = load_llff_data(args.datadir, args.factor, \n",
    "                                                          recenter=True, bd_factor=.75, \n",
    "                                                          spherify=args.spherify)\n",
    "# images 20 756 1008 3\n",
    "# poses 20 3 5\n",
    "# render_poses 120 3 5 # 120 是渲染的数量\n",
    "H, W, focal = poses[0,:3,-1].astype(np.float32)\n",
    "\n",
    "H = int(H)\n",
    "W = int(W)\n",
    "hwf = [H, W, focal]\n",
    "\n",
    "images = images.astype(np.float32)\n",
    "poses = poses.astype(np.float32)\n",
    "\n",
    "if args.no_ndc:\n",
    "    near = tf.reduce_min(bds) * .9\n",
    "    far = tf.reduce_max(bds) * 1.\n",
    "else:\n",
    "    near = 0.\n",
    "    far = 1.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create nerf model\n",
    "_, render_kwargs_test, start, grad_vars, models = run_nerf.create_nerf(args)\n",
    "\n",
    "bds_dict = {\n",
    "    'near' : tf.cast(near, tf.float32),\n",
    "    'far' : tf.cast(far, tf.float32),\n",
    "}\n",
    "render_kwargs_test.update(bds_dict)\n",
    "\n",
    "print('Render kwargs:')\n",
    "pprint.pprint(render_kwargs_test)\n",
    "\n",
    "\n",
    "down = 4\n",
    "render_kwargs_fast = {k : render_kwargs_test[k] for k in render_kwargs_test}\n",
    "render_kwargs_fast['N_importance'] = 0\n",
    "\n",
    "c2w = np.eye(4)[:3,:4].astype(np.float32) # identity pose matrix\n",
    "test = run_nerf.render(H//down, W//down, focal/down, c2w=c2w, **render_kwargs_fast)\n",
    "img = np.clip(test[0],0,1)\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "down = 8 # trade off resolution+aliasing for render speed to make this video faster\n",
    "frames = []\n",
    "for i, c2w in enumerate(render_poses):\n",
    "    # c2w 3 5 \n",
    "    if i%8==0: print(i)\n",
    "    test = run_nerf.render(H//down, W//down, focal/down, c2w=c2w[:3,:4], **render_kwargs_fast)\n",
    "    frames.append((255*np.clip(test[0],0,1)).astype(np.uint8))\n",
    "    \n",
    "print('done, saving')\n",
    "f = 'logs/fern_example/video.mp4'\n",
    "imageio.mimwrite(f, frames, fps=30, quality=8)\n",
    "\n",
    "from IPython.display import Video\n",
    "Video(f, height=320)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from ipywidgets import interactive, widgets\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def f(x, y, z):\n",
    "    \n",
    "    c2w = tf.convert_to_tensor([\n",
    "        [1,0,0,x],\n",
    "        [0,1,0,y],\n",
    "        [0,0,1,z],\n",
    "        [0,0,0,1],\n",
    "    ], dtype=tf.float32)\n",
    "    \n",
    "    test = run_nerf.render(H//down, W//down, focal/down, c2w=c2w, **render_kwargs_fast)\n",
    "    img = np.clip(test[0],0,1)\n",
    "    \n",
    "    plt.figure(2, figsize=(20,6))\n",
    "    plt.imshow(img)\n",
    "    plt.show()\n",
    "    \n",
    "\n",
    "sldr = lambda : widgets.FloatSlider(\n",
    "    value=0.,\n",
    "    min=-1.,\n",
    "    max=1.,\n",
    "    step=.01,\n",
    ")\n",
    "\n",
    "names = ['x', 'y', 'z']\n",
    "    \n",
    "interactive_plot = interactive(f, **{n : sldr() for n in names})\n",
    "interactive_plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7bcd4bc356ddf5f625d74850a8c16eeecb3909da86156df6e11fed4fdd37981"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
